from keras import backend as K
from tensorflow.examples.tutorials.mnist import input_data
from keras.models import Model
from keras.layers import Input, Conv2D, MaxPooling2D, Reshape, Lambda
from CapsuleLayer import Capsule, margin_loss


# 载入数据
def read_data(path):
    mnist = input_data.read_data_sets(path, one_hot=True)
    train_x, train_y = mnist.train.images.reshape(-1, 28, 28, 1), mnist.train.labels,
    valid_x, valid_y = mnist.validation.images.reshape(-1, 28, 28, 1), mnist.validation.labels,
    test_x, test_y = mnist.test.images.reshape(-1, 28, 28, 1), mnist.test.labels
    return train_x, train_y, valid_x, valid_y, test_x, test_y


# 模型
def MODEL():
    inputs = Input(shape=(28, 28, 1))
    x = Conv2D(16, (5, 5), padding='same', activation='relu')(inputs)
    x = MaxPooling2D(pool_size=(2, 2))(x)
    x = Conv2D(32, (5, 5), padding='same', activation='relu')(x)
    x = MaxPooling2D(pool_size=(2, 2))(x)
    x = Reshape((-1, 32))(x)  # [None, 49, 32] 即前一层胶囊 [None, input_num, input_dim]
    x = Capsule(num_capsule=10, dim_capsule=30, routings=5)(x)
    output = Lambda(lambda z: K.sqrt(K.sum(K.square(z), axis=2)))(x)  # 每个胶囊取模长
    model = Model(inputs=inputs, output=output)
    return model


# 主函数
def main(train_x, train_y, valid_x, valid_y, test_x, test_y):
    model = MODEL()
    model.compile(loss=margin_loss, optimizer='adam', metrics=['accuracy'])
    model.summary()
    model.fit(train_x, train_y, batch_size=500, nb_epoch=20, verbose=2)
    pre = model.evaluate(test_x, test_y, batch_size=500, verbose=2)
    print('test_loss:', pre[0], '- test_acc:', pre[1])


train_x, train_y, valid_x, valid_y, test_x, test_y = read_data('MNIST_data')
main(train_x, train_y, valid_x, valid_y, test_x, test_y)
